
/*
优化谱减法降噪程序
*/
#include "fftwrap.h"
#include "arch.h"
#include "speex/speex_preprocess.h"
#include "speex/speex_echo.h"
#include "os_support.h"
#include "math_approx.h"
#include <math.h>
#include <stdint.h>

#define EXPORT

#define MIN(a,b) ((a) < (b) ? (a) : (b))
#define MAX(a,b) ((a) > (b) ? (a) : (b))

static void SqrtHanning(float *win, int len)
{
    int n; float PI = 3.1415926;
	for (n=0; n<len/2+1; n++)
	{
		win[n] = sqrtf(0.5*(1-cos(2*PI*n/len)));
	}
	for (n=1; n<len/2; n++)
	{
		win[len/2+n] = win[len/2-n];
	}
}

// add window
static void window_data(float *in, float *out, float *win, int len)
{
    int i;
    for (i = 0; i < len; i++)
    {
        out[i] = in[i] * win[i];
    }
}

struct SpeexPreprocessState_
{
    void  *fft_lookup;
    int frame_size;
    int window_size;
    float *win;
    float *winbuf;
    float *overlap;
    float *E; // echo
    float *e;
    //
    float *ps;
    int nb_adapt;
    float *old_ps;
    float *noise;

    float *S;
    float *Smin;
    int min_count;
    float *Stmp;
    float *update_prob;
    
    float *gain;
    float *gain2;
    float *gain_floor;
    float *post;
    float *prior;
    float *zeta;
    /**/
    SpeexEchoState *echo_state;
    spx_word32_t *echo_noise;
    spx_word32_t *residual_echo;
};


EXPORT SpeexPreprocessState *speex_preprocess_state_init(int frame_size, int sampling_rate)
{
    SpeexPreprocessState *self = (SpeexPreprocessState *)speex_alloc(sizeof(SpeexPreprocessState));
    self->frame_size = frame_size;
    self->window_size = 2*frame_size;
    int N = 2*frame_size;
    self->fft_lookup = spx_fft_init(N);
    /**/
    self->win = (spx_word16_t*)speex_alloc(N*sizeof(spx_word16_t));
    SqrtHanning(self->win, N);
    self->overlap = (spx_word16_t*)speex_alloc(self->frame_size*sizeof(spx_word16_t));
    for (int n=0; n<self->frame_size; n++) {
        self->overlap[n] = 0.0f;
    }
    self->winbuf = (spx_word16_t*)speex_alloc(N*sizeof(spx_word16_t));
    /**/
    self->e = (spx_word16_t*)speex_alloc(N*sizeof(spx_word16_t));
    self->E = (spx_word16_t*)speex_alloc(N*sizeof(spx_word16_t));
    //
    self->noise = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->ps = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->old_ps = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->Smin = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->S = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->Stmp = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->update_prob = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    
    self->gain = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->gain2 = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->gain_floor = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));

    self->post = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->prior = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->zeta = (spx_word32_t *)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->min_count = 0;
    self->nb_adapt = 0;

    for (int nf=0; nf<self->frame_size; nf++) {
        self->noise[nf] = 1.0f;
        self->old_ps[nf] = 1.0f;
        self->gain[nf] = 1.0f;
        self->post[nf] = 1.0f;
        self->prior[nf] = 1.0f;
        self->update_prob[nf] = 1.0f;
    }
    /**/
    self->echo_noise = (spx_word32_t*)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    self->residual_echo = (spx_word32_t*)speex_alloc((self->frame_size)*sizeof(spx_word32_t));
    return self;
}

EXPORT void speex_preprocess_state_destroy(SpeexPreprocessState *st)
{
    spx_fft_destroy(st->fft_lookup);
    speex_free(st->win);
    speex_free(st->overlap);
    speex_free(st->winbuf);
    speex_free(st->E);
    speex_free(st->e);
    //
    speex_free(st->noise);
    speex_free(st->ps);
    speex_free(st->old_ps);
    speex_free(st->S);
    speex_free(st->Smin);
    speex_free(st->Stmp);
    speex_free(st->update_prob);
    speex_free(st->gain);
    speex_free(st->gain2);
    speex_free(st->gain_floor);
    speex_free(st->post);
    speex_free(st->prior);
    speex_free(st->zeta);
    /**/
    speex_free(st->echo_noise);
    speex_free(st->residual_echo);

    speex_free(st);
}

void speex_echo_get_residual(SpeexEchoState *st, spx_word32_t *Yout, int len);

EXPORT int speex_preprocess_run(SpeexPreprocessState *st, spx_int16_t *x)
{
    SpeexPreprocessState *self = st;
    int N = self->window_size;
    memcpy(self->e, self->e+self->frame_size, sizeof(spx_word16_t)*(self->frame_size));    
    for (int n=0; n<self->frame_size; n++) {
        self->e[self->frame_size + n] = (float)x[n];
    }
    window_data(self->e, self->winbuf, self->win, N);
    
    spx_fft(self->fft_lookup, self->winbuf, self->E);

    if (st->echo_state) {
        speex_echo_get_residual(st->echo_state, st->residual_echo, self->frame_size);
        for (int i=0; i<self->frame_size; i++) {
            st->echo_noise[i] = 0.0f; // MAX32(MULT16_32_Q15(QCONST16(.45f,15),st->echo_noise[i]), st->residual_echo[i]);
        }
    }
    else {
        for (int i=0; i<self->frame_size; i++) {
            st->echo_noise[i] = 0.0f;
        }
    }
    
    spx_word16_t beta, beta_1;
    self->nb_adapt++;
    if (self->nb_adapt>20000) {
        self->nb_adapt = 20000;
    }
    self->min_count++;
    beta = MAX(0.03f, 1/self->nb_adapt);
    beta_1 = 1.f - beta;

    // if (/*echo*/) {
    //     // 
    // }
    // preprocess_analysis
    self->ps[0] = self->E[0]*self->E[0];
    for (int nf=1; nf<self->frame_size; nf++) {
        self->ps[nf] = self->E[2*nf-1]*self->E[2*nf-1] + self->E[2*nf]*self->E[2*nf];
    }
    // update_noise_prob
    for (int nf=1; nf<self->frame_size-1; nf++)
        self->S[nf] = 0.8*self->S[nf] + 0.05*self->ps[nf-1] + 0.1*self->ps[nf] + 0.05*self->ps[nf+1];
    self->S[0] = 0.8*self->S[0] + 0.2*self->ps[0];
    self->S[self->frame_size-1] = 0.8*self->S[self->frame_size-1] + 0.2*self->ps[self->frame_size-1];
    if (self->nb_adapt==1)
    {
        for (int nf=0; nf<self->frame_size; nf++)
            self->Smin[nf] = self->Stmp[nf] = 0;
    }
    int min_range;
    if (self->nb_adapt < 100)
        min_range = 15;
    else if (self->nb_adapt < 1000)
        min_range = 50;
    else if (self->nb_adapt < 10000)
        min_range = 150;
    else
        min_range = 300;
    if (self->min_count > min_range)
    {
        self->min_count = 0;
        for (int nf=0; nf<self->frame_size; nf++)
        {
            self->Smin[nf] = MIN(self->Stmp[nf], self->S[nf]);
            self->Stmp[nf] = self->S[nf];
        }
    } else {
        for (int nf=0; nf<self->frame_size; nf++)
        {
            self->Smin[nf] = MIN(self->Smin[nf], self->S[nf]);
            self->Stmp[nf] = MIN(self->Stmp[nf], self->S[nf]);      
        }
    }
    for (int nf=0; nf<self->frame_size; nf++) {
        if (0.4*self->S[nf] > self->Smin[nf]) {
            self->update_prob[nf] = 1;
        }
        else {
            self->update_prob[nf] = 0;
        }
    }
    // Update the noise estimate for the frequencies where it can be
    for (int i=0; i<self->frame_size; i++) {
        if (!self->update_prob[i] || self->ps[i] < self->noise[i]) {
            self->noise[i] = MAX(0, beta_1*self->noise[i]) + beta*self->ps[i]; // X(jw)^2
        }
    }
    if (self->nb_adapt == 1) {
        for (int i=0; i>self->frame_size; i++) {
            self->old_ps[i] = self->ps[i];
        }
    }

    #ifdef LOGMMSE
    float ksi; float gamma_k; float ksi_min = pow(10,-2.5);
    for (int i=1; i<self->frame_size; i++) {
        gamma_k = MIN(self->ps[i]/self->noise[i], 40);
        if (self->nb_adapt==1) {
            ksi = 0.98 + 0.02*MAX(gamma_k-1, 0);
        }
        else {
            ksi = 0.98*self->old_ps[i]/self->noise[i] + 0.02*MAX(gamma_k-1, 0);
            ksi = MAX(ksi_min, ksi);
        }
        float A = ksi/(1+ksi);
        float vk = A*gamma_k;
        float ei_vk = 0.5*(-expp_approx(vk));
        self->gain2[i] = A*expf_approx(ei_vk);
        // printf("%.3f\n", self->gain2[i]);
    }
    #else 
    // spectral subtraction
    float alpha = 6; float alpha2 = 0.05; // 6或者3, 0.1或者0.01
    for (int i=1; i<self->frame_size; i++) {
        if (self->ps[i]-(alpha+alpha2)*(self->noise[i]+self->echo_noise[i])>0)
            self->gain2[i] = sqrt_approx((self->ps[i]-alpha*(self->noise[i]+self->echo_noise[i]))/self->ps[i]);
        else 
            self->gain2[i] = sqrt_approx(alpha2*self->noise[i]/self->ps[i]);
    }
    #endif    

    self->E[0] = 0.0f;
    for (int nf=0; nf<self->frame_size; nf++) {
        self->E[2*nf-1] = self->gain2[nf]*self->E[2*nf-1];
        self->E[2*nf] = self->gain2[nf]*self->E[2*nf];
        #ifdef LOGMMSE
        self->old_ps[nf] = self->E[2*nf-1]*self->E[2*nf-1] + self->E[2*nf]*self->E[2*nf];
        #endif
    }
    self->E[2*self->frame_size-1] = 0.0f;
    
    spx_ifft(self->fft_lookup, self->E, self->winbuf);
     
    float tmp;
    for (int n=0; n<self->frame_size; n++) {
        tmp = self->winbuf[n]*self->win[n] + self->overlap[n];
        x[n] = (short)tmp;
    }
    for (int n=self->frame_size; n<N; n++) {
        self->overlap[n-self->frame_size] = self->winbuf[n]*self->win[n];
    }

    return 0;
}

EXPORT int speex_preprocess_ctl(SpeexPreprocessState *state, int request, void *ptr)
{
    SpeexPreprocessState *st;
    st = (SpeexPreprocessState*)state;
    switch (request)
    {
    case SPEEX_PREPROCESS_SET_ECHO_STATE: {
        st->echo_state = (SpeexEchoState*)ptr;
    }
    break;
    case SPEEX_PREPROCESS_GET_ECHO_STATE: {
        (*(SpeexEchoState**)ptr) = (SpeexEchoState*)st->echo_state;
    }
    break;    
    default:
        break;
    }
}